-- | Module:        Hexena.Internal.Blooming
--   Stability:     Experimental
--
-- Internal-Use module.  This module exports variables and functions
-- that are considered not for public consumption.  As such with all
-- internal-use modules.  The variables contained are considered to be
-- in flux and may be changed at any time as the project unfolds.
-- Developers use contents contained herein at your own risk.

module Hexena.Internal.Blooming ( HashFun
                                , Blooming (..)
                                , MBlooming (..)
                                , bloomingNewM
                                , bloomingInsertM
                                , bloomingElem
                                , bloomingElemM
                                , bloomingFreezeM
                                , bloomingQueryFillM
                                , bloomingSerialiseM
                                , createUniqueNGrams
                                , createUniqueNGrams'
                                , defaultHashes
                                , doHashes
                                , fillBloomFilter'
                                , fillBloomFromFile
                                , fillBloomFromFileIO
                                , loadBLOM
                                , saveBLOM
                                , takeN
                                , takeN'
                                ) where

import System.IO (stdout)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.Set as S
import Data.Array.IO
import Data.Array.Unboxed ((!), UArray, elems, listArray)
import Data.Binary.Get
import Data.Binary.Put
import Data.Bits ((.|.), (.&.), bit, shiftR)
import Data.BloomFilter
import Data.IFF
import Data.List (find, foldl')
import Data.Maybe
import Data.Word

import Hexena.Chunks

type HashFun = [Word8] -> Hash

-- | Data structure for a bloom filter in the IO Monad.
--
-- This follows mostly the code found in the hackage bloomfilter
-- package.
data MBlooming a = MBlooming { bloomingHashesM :: !(a -> [Hash])
                             , bloomingModuloM :: Word32
                             , bloomingBitArrayM :: !(IOUArray Word32 Word8)
                             }

data Blooming a = Blooming { bloomingHashes :: !(a -> [Hash])
                           , bloomingModulo :: Word32
                           , bloomingBitArray :: !(UArray Word32 Word8)
                           }

loadBLOM :: (a -> [Hash]) -> T -> Maybe (Blooming a)
loadBLOM hashF iff = if blomID /= chunkId iff then Nothing else Just $ Blooming hashF modulo arr
    where formBLOM = parts $ chunk $ iff :: [T]
          chunkMOD = content $ chunk $ fromJust $ find ((==) (chunkIdFromString "MOD ") . chunkId) formBLOM
          chunkBARR = content $ chunk $ fromJust $ find ((==) (chunkIdFromString "BARR") . chunkId) formBLOM
          modulo = runGet (getWord32be) $ BL.fromChunks [chunkMOD]
          arr = listArray (0, modulo) $ B.unpack chunkBARR
          blomID = chunkIdFromString "BLOM"

saveBLOM :: Blooming a -> T
saveBLOM bf = Cons (chunkIdFromString "BLOM") $ Form [cMOD, cBARR]
    where modulo = putWord32be $ bloomingModulo bf
          cMOD = Cons (chunkIdFromString "MOD ") $ Chunk $ B.concat $ BL.toChunks $ runPut modulo
          cBARR = Cons (chunkIdFromString "BARR") $ Chunk $ B.pack $ elems $ bloomingBitArray bf

bloomingFreezeM :: MBlooming a -> IO (Blooming a)
bloomingFreezeM bf = Blooming hf m `fmap` freeze barr
    where hf = bloomingHashesM bf
          m = bloomingModuloM bf
          barr = bloomingBitArrayM bf

bloomingSerialiseM :: MBlooming a -> IO (Word32, UArray Word32 Word8)
bloomingSerialiseM bf = do arr <- freeze $ bloomingBitArrayM bf
                           return (bloomingModuloM bf, arr)

bloomingCheckBit :: Blooming a -> Hash -> Bool
bloomingCheckBit mb h = let h' = mod h $ bloomingModulo mb
                            idx = h' `shiftR` 3
                            theBit = bit (fromIntegral $ h' .&. 7)
                            arr = bloomingBitArray mb
                        in arr ! idx .&. theBit /= 0

bloomingElem :: Blooming a -> a -> Bool
bloomingElem mb ele = all id $ map (bloomingCheckBit mb) hashes
  where hashes = bloomingHashes mb ele



-- | Create a new mutable Bloom filter. 
--
bloomingNewM :: (a -> [Hash])    -- ^ family of hash functions to use
             -> Word32           -- ^ number of bits in filter
             -> IO (MBlooming a)
bloomingNewM hashes modulo = do arr <- newArray (0, modulo `div` 8) 0
                                return $ MBlooming hashes modulo arr

bloomingCheckBitM :: MBlooming a -> Hash -> IO Bool
bloomingCheckBitM mb h = let h' = mod h $ bloomingModuloM mb
                             idx = h' `shiftR` 3
                             theBit = bit (fromIntegral $ h' .&. 7)
                             arr = bloomingBitArrayM mb
                         in do x <- readArray arr idx
                               return $ x .&. theBit /= 0

bloomingQueryFillM :: MBlooming a -> IO Word32
bloomingQueryFillM mb = (sum . map (f . g)) `fmap` getElems arr
    where arr = bloomingBitArrayM mb
          g :: Word8 -> Word64
          g x = (fromIntegral x * 0x808040201) .&. 0x111111111
          f :: Word64 -> Word32
          f = fromIntegral . (`mod` 0xf)

bloomingSetBitM :: MBlooming a -> Hash -> IO ()
bloomingSetBitM mb h = let h' = mod h $ bloomingModuloM mb
                           idx = h' `shiftR` 3
                           theBit = bit (fromIntegral $ h' .&. 7)
                           arr = bloomingBitArrayM mb
                       in do x <- readArray arr idx
                             writeArray arr idx $ x .|. theBit

bloomingInsertM :: MBlooming a -> a -> IO ()
bloomingInsertM mb = mapM_ (bloomingSetBitM mb) . bloomingHashesM mb

bloomingElemM :: MBlooming a -> a -> IO Bool
bloomingElemM mb ele = all id `fmap` mapM (bloomingCheckBitM mb) hashes
  where hashes = bloomingHashesM mb ele

defaultHashes :: [HashFun]
defaultHashes = [ash, jsHash, sdbmHash]

ash :: [Word8] -> Hash
ash = fromIntegral . flip mod 0x7fffffff . f
    where f :: [Word8] -> Word32
          f = foldl' (\a b -> a `seq` a * 13 + fromIntegral b) 7

jsHash :: [Word8] -> Hash
jsHash = fromIntegral . flip mod 0x7fffffff . f
    where f :: [Word8] -> Word32
          f = foldl' (\a b -> a `seq` (a * 32 + a `div` 4 + fromIntegral b)) 1315423911

sdbmHash :: [Word8] -> Hash
sdbmHash = fromIntegral . flip mod 0x7fffffff . f
    where f :: [Word8] -> Word32
          f = foldl' (\a b -> a `seq` (a * 64 + a * 65536 - a + fromIntegral b)) 0

takeN :: Int -> BL.ByteString -> [BL.ByteString]
takeN n xs | BL.length xs >= fromIntegral n = BL.take (fromIntegral n) xs : takeN n (BL.drop 1 xs)
           | otherwise = [xs]

takeN' :: Int -> B.ByteString -> [B.ByteString]
takeN' n xs = if B.length xs < n' then [] else B.take n' xs : takeN' n (B.drop 1 xs)
    where n' = fromIntegral n

out :: S.Set BL.ByteString -> IO ()
out = mapM_ (BL.hPutStr stdout) . S.toList

doHashes :: [Word8] -> [HashFun] -> [Hash]
doHashes bs = map (\x -> x bs)

createUniqueNGrams :: Int -> BL.ByteString -> [BL.ByteString]
createUniqueNGrams n = S.toList . S.fromList . takeN n

createUniqueNGrams' :: Int -> B.ByteString -> [B.ByteString]
createUniqueNGrams' n = S.toList . S.fromList . takeN' n

hashNGrams :: [BL.ByteString] -> [Hash]
hashNGrams = concatMap (flip doHashes [ash, jsHash, sdbmHash] . BL.unpack)

-- | Fill a default Bloom filter with maximum n elements from ByteStrings
fillBloomFilter' :: Int -> [B.ByteString] -> Bloom B.ByteString
fillBloomFilter' = fromListB (flip doHashes defaultHashes . B.unpack)

-- | Fill a Bloom (maximum n elements) filter with m-grams of size m.
fillBloomFromFile :: Int -> Int -> [FilePath] -> IO (Bloom B.ByteString)
fillBloomFromFile m n fns = do mgrams <- concatMap (takeN' m) `fmap` mapM B.readFile fns
                               return $ fillBloomFilter' n mgrams

fillBloomFromFileIO :: Int -> Word32 -> [FilePath] -> IO (MBlooming B.ByteString)
fillBloomFromFileIO m n fns = do mgrams <- concatMap (takeN' m) `fmap` mapM B.readFile fns
                                 bf <- bloomingNewM (flip doHashes defaultHashes . B.unpack) n
                                 mapM_ (bloomingInsertM bf) mgrams
                                 return bf


-- all_filenames <- (Prelude.map (path_prefix ++) . Prelude.filter ((".o" ==) . takeExtension)) `fmap` getDirectoryContents path_prefix
-- bf <- fillBloomFromFile 4 1000000 all_filenames
-- (\x -> let y = takeN' 4 x in (Prelude.length y, Prelude.length $ Prelude.filter id $ Prelude.map (Prelude.flip elemB bf) y)) `fmap` Data.ByteString.readFile ""

-- # -- | Hash all files hashing all ngrams... We do not need to uniquelify
-- # -- | the list of hashes?! Putting it in the bloomap will do that for
-- # -- | us. (Is this really faster? Yes it seems so...)
-- # main2 :: [BL.ByteString] -> IO ()
-- # main2 xs = let hashes = concatMap (hashNGrams . takeN 5) xs
-- #                bloomap = hashes `seq` IS.fromList $ map (flip mod 0xFFFFFFF) hashes
-- #                bloosize = bloomap `seq` IS.size bloomap
-- #            in do print $ length hashes
-- #                  putStrLn $ show bloosize ++ "\t" ++ show bloomap
-- #

